import os
import json
from queue import Queue
from tree_sitter_languages import get_parser
import argparse
'''==============================SEAST解析 处理源代码为SEAST节点序列=================================='''
class ASTParser:
    def __init__(self):
        # 初始化解析器
        self.parser = None
        self.current_language = None

    def extract_node_value_or_type(self, node):
        """
        提取节点的值或类型，如果是变量值节点则提取其值
        输入:节点
        输出:节点的值或类型
        """
        # 对于常量和变量值，直接提取其值
        if node.type in ['string', 'number', 'identifier']:
            value = self._get_node_value(node)
            return value
        # 其他节点则提取类型
        else:
            return node.type

    def _get_node_value(self, node):
        """
        提取节点的值
        输入:节点
        输出:节点的值
        """
        if node.type == 'string':
            return node.text.decode('gbk')  # 处理字符串类型的值
        elif node.type == 'number':
            return node.text.decode('gbk')  # 处理数字类型的值
        elif node.type == 'identifier':
            return node.text.decode('gbk')  # 处理标识符（变量名）
        return None

    def extract_node_type_sequence(self, root):
        """
        提取AST节点类型序列, 对于变量值节点提取其值
        输入:AST根节点
        输出:节点类型序列
        """
        q = Queue()
        node_type_sequence = []
        q.put(root)
        while not q.empty():
            node = q.get()
            for child in node.children:
                q.put(child)
                node_value_or_type = self.extract_node_value_or_type(child)
                node_type_sequence.append(node_value_or_type)
        return node_type_sequence

    def load_existing_entries(self, output_file):
        """
        从output_file中读取已有的条目
        输入:输出文件路径
        输出:已有的条目集合
        """
        existing_entries = set()
        if os.path.exists(output_file):
            with open(output_file, 'r') as infile:
                for line in infile:
                    entry = json.loads(line)
                    existing_entries.add(entry['name'])
        return existing_entries

    def process_dataset(self, dataset_path, output_file):
        """
        处理数据集
        输入:数据集路径,输出文件路径
        输出:无
        """
        # 按语言类型统计
        stats = {
            'py': {'write': 0, 'fail': 0},
            'js': {'write': 0, 'fail': 0},
            'perl': {'write': 0, 'fail': 0},
            'php': {'write': 0, 'fail': 0}
        }
        total_write = 0

        # 加载现有的条目
        existing_entries = self.load_existing_entries(output_file)

        for root, _, files in os.walk(dataset_path):
            for file in files:
                file_path = os.path.join(root, file)
                file_ext = os.path.splitext(file)[1]
                
                # 根据文件类型选择对应的语言解析器
                if file_ext == '.py':
                    self.parser = get_parser('python')
                    lang = 'py'
                elif file_ext == '.js':
                    self.parser = get_parser('javascript')
                    lang = 'js'
                elif file_ext == '.pl':
                    self.parser = get_parser('perl')
                    lang = 'perl'
                elif file_ext == '.php':
                    self.parser = get_parser('php')
                    lang = 'php'
                else:
                    continue

                # 根据目录路径判断标签
                if "beg" in file_path.lower():
                    label = 0
                elif "mal" in file_path.lower():
                    label = 1
                else:
                    continue

                if file in existing_entries:
                    continue
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        code = f.read()

                    # 解析文件并提取节点类型序列
                    tree = self.parser.parse(bytes(code, 'utf-8'))
                    root_node = tree.root_node
                    node_type_sequence = self.extract_node_type_sequence(root_node)
                    
                    entry = {
                        'name': file,
                        'features': node_type_sequence,
                        'language': lang,
                        'label': label
                    }
                    with open(output_file, 'a') as outfile:
                        json.dump(entry, outfile)
                        outfile.write('\n')
                    stats[lang]['write'] += 1
                    total_write += 1
                    
                    # 每1000个文件打印一次统计信息
                    if total_write % 2000 == 0:
                        print(f"\n当前处理进度: ")
                        for l in ['py', 'js', 'perl', 'php']:
                            print(f"{l.upper()}: 成功写入 {stats[l]['write']} 个文件, 失败 {stats[l]['fail']} 个文件")
                        print(f"总共写入: {total_write} 个文件\n")
                    
                    # 添加到existing_entries中
                    existing_entries.add(file)
                    
                    # 释放tree和code的引用，帮助垃圾回收
                    del tree
                    del code
                except Exception:
                    try:
                        with open(file_path, 'r', encoding='gbk') as f:
                            code = f.read()

                        # 解析文件并提取节点类型序列
                        tree = self.parser.parse(bytes(code, 'gbk'))
                        root_node = tree.root_node
                        node_type_sequence = self.extract_node_type_sequence(root_node)
                        
                        entry = {
                            'name': file,
                            'features': node_type_sequence,
                            'language': lang,
                            'label': label
                        }
                        with open(output_file, 'a') as outfile:
                            json.dump(entry, outfile)
                            outfile.write('\n')
                        stats[lang]['write'] += 1
                        total_write += 1
                    except Exception:
                        stats[lang]['fail'] += 1
                        # 不打印错误内容
                        pass
        print("\n处理完成最终统计: ")
        for lang in ['py', 'js', 'perl', 'php']:
            print(f"{lang.upper()}: 成功写入 {stats[lang]['write']} 个文件, 失败 {stats[lang]['fail']} 个文件")
        print(f"总共写入: {total_write} 个文件")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='SEAST解析器')
    parser.add_argument('--data_dir', default='../dataraw', help='输入数据集路径')
    parser.add_argument('--output_file', default='./seast_text.json', help='输出文件路径')
    
    args = parser.parse_args()
    
    ast_parser = ASTParser()
    print("开始处理数据集...")
    ast_parser.process_dataset(args.data_dir, args.output_file)
    print(f"处理完成，结果已写入{args.output_file}")